import numpy as np
from scipy.special import jv, yv

'''
需要 numpy、scipy
nMedium 是介质折射率  默认值为1
asDict 是否按照字典输出
asCrossSection 是否输出消光截面、散射截面、吸收截面
nMedium, asDict, asCrossSection 可不设置 例：MieQ(1.77 + 0.63j, 375, 300) ，不设置会使用默认值
这段程序是在 PyMieScatt(一个开源的Mie散射计算包) 中提取的部分代码
'''


def MieQ(m, wavelength, diameter, nMedium=1.0, asDict=False, asCrossSection=False):
    #  米散射计算
    nMedium = nMedium.real  # 介质折射率
    m /= nMedium
    wavelength /= nMedium
    x = np.pi * diameter / wavelength  # 粒子尺度参数
    if x == 0:
        return 0, 0, 0, 1.5, 0, 0, 0
    elif x <= 0.05:  # 尺度参数小于0.05 可用瑞利散射计算
        return RayleighMieQ(m, wavelength, diameter, nMedium, asDict)
    elif x > 0.05:  # 尺度参数大于0.05 计算Mie散射
        nmax = np.round(2 + x + 4 * (x ** (1 / 3)))  # n的最大项
        n = np.arange(1, nmax + 1)
        n1 = 2 * n + 1
        n2 = n * (n + 2) / (n + 1)
        n3 = n1 / (n * (n + 1))
        x2 = x ** 2

        an, bn = Mie_ab(m, x)  # an、bn计算

        # 消光系数、散射系数、吸收系数
        qext = (2 / x2) * np.sum(n1 * (an.real + bn.real))
        qsca = (2 / x2) * np.sum(n1 * (an.real ** 2 + an.imag ** 2 + bn.real ** 2 + bn.imag ** 2))
        qabs = qext - qsca

        g1 = [an.real[1:int(nmax)],
              an.imag[1:int(nmax)],
              bn.real[1:int(nmax)],
              bn.imag[1:int(nmax)]]
        g1 = [np.append(x, 0.0) for x in g1]
        # 不对称因子
        g = (4 / (qsca * x2)) * np.sum(
            (n2 * (an.real * g1[0] + an.imag * g1[1] + bn.real * g1[2] + bn.imag * g1[3])) + (
                    n3 * (an.real * bn.real + an.imag * bn.imag)))

        qpr = qext - qsca * g
        # 后向散射系数
        qback = (1 / x2) * (np.abs(np.sum(n1 * ((-1) ** n) * (an - bn))) ** 2)
        qratio = qback / qsca
        if asCrossSection:
            css = np.pi * (diameter / 2) ** 2
            cext = css * qext
            csca = css * qsca
            cabs = css * qabs
            cpr = css * qpr
            cback = css * qback
            cratio = css * qratio
            if asDict:
                return dict(Cext=cext, Csca=csca, Cabs=cabs, g=g, Cpr=cpr, Cback=cback, Cratio=cratio)
            else:
                return cext, csca, cabs, g, cpr, cback, cratio
        else:
            if asDict:
                return dict(Qext=qext, Qsca=qsca, Qabs=qabs, g=g, Qpr=qpr, Qback=qback, Qratio=qratio)
            else:
                return qext, qsca, qabs, g, qpr, qback, qratio


def Mie_ab(m, x):
    #  米散射的an、bn
    mx = m * x
    nmax = np.round(2 + x + 4 * (x ** (1 / 3)))
    nmx = np.round(max(nmax, np.abs(mx)) + 16)
    n = np.arange(1, nmax + 1)  #
    nu = n + 0.5  #

    sx = np.sqrt(0.5 * np.pi * x)

    px = sx * jv(nu, x)  #
    p1x = np.append(np.sin(x), px[0:int(nmax) - 1])  #

    chx = -sx * yv(nu, x)  #
    ch1x = np.append(np.cos(x), chx[0:int(nmax) - 1])  #

    gsx = px - (0 + 1j) * chx  #
    gs1x = p1x - (0 + 1j) * ch1x  #

    # B&H Equation 4.89
    Dn = np.zeros(int(nmx), dtype=complex)
    for i in range(int(nmx) - 1, 1, -1):
        Dn[i - 1] = (i / mx) - (1 / (Dn[i] + i / mx))

    D = Dn[1:int(nmax) + 1]  # Dn(mx), 有限个项
    da = D / m + n / x
    db = m * D + n / x

    an = (da * px - p1x) / (da * gsx - gs1x)
    bn = (db * px - p1x) / (db * gsx - gs1x)

    return an, bn


def RayleighMieQ(m, wavelength, diameter, nMedium=1.0, asDict=False, asCrossSection=False):
    #  瑞利散射
    nMedium = nMedium.real
    m /= nMedium
    wavelength /= nMedium
    x = np.pi * diameter / wavelength
    if x == 0:
        return 0, 0, 0, 1.5, 0, 0, 0
    elif x > 0:
        LL = (m ** 2 - 1) / (m ** 2 + 2)  # Lorentz-Lorenz term
        LLabsSq = np.abs(LL) ** 2
        qsca = 8 * LLabsSq * (x ** 4) / 3
        qabs = 4 * x * LL.imag
        qext = qsca + qabs
        qback = 1.5 * qsca
        qratio = 1.5
        g = 0
        qpr = qext
        if asCrossSection:
            css = np.pi * (diameter / 2) ** 2
            cext = css * qext
            csca = css * qsca
            cabs = css * qabs
            cpr = css * qpr
            cback = css * qback
            cratio = css * qratio
            if asDict:
                return dict(Cext=cext, Csca=csca, Cabs=cabs, g=g, Cpr=cpr, Cback=cback, Cratio=cratio)
            else:
                return cext, csca, cabs, g, cpr, cback, cratio
        else:
            if asDict:
                return dict(Qext=qext, Qsca=qsca, Qabs=qabs, g=g, Qpr=qpr, Qback=qback, Qratio=qratio)
            else:
                return qext, qsca, qabs, g, qpr, qback, qratio


def MieQCoreShell(mCore, mShell, wavelength, dCore, dShell, nMedium=1.0, asDict=False, asCrossSection=False):

    xCore = np.pi * dCore / wavelength
    xShell = np.pi * dShell / wavelength
    # 这些时候是Mie散射计算
    if xCore == xShell:
        return MieQ(mCore, wavelength, dShell, nMedium=nMedium, asDict=asDict, asCrossSection=asCrossSection)
    elif xCore == 0:
        return MieQ(mShell, wavelength, dShell, nMedium=nMedium, asDict=asDict, asCrossSection=asCrossSection)
    elif mCore == mShell:
        return MieQ(mCore, wavelength, dShell, nMedium=nMedium, asDict=asDict, asCrossSection=asCrossSection)
    elif xCore > 0:
        nMedium = nMedium.real
        # wavelength /= nMedium  # 要么选择重新定义波长，要么xCore和xShell  这里重新定义xCore和xShell
        mCore /= nMedium
        mShell /= nMedium
        xCore = np.pi * dCore * nMedium / wavelength  # 重新定义xCore和xShell并不理想，但似乎有必要保持MieQ条件不变
        xShell = np.pi * dShell * nMedium / wavelength

        nmax = np.round(2 + xShell + 4 * (xShell ** (1 / 3)))
        n = np.arange(1, nmax + 1)
        n1 = 2 * n + 1
        n2 = n * (n + 2) / (n + 1)
        n3 = n1 / (n * (n + 1))
        xShell2 = xShell ** 2

        an, bn = CoreShell_ab(mCore, mShell, xCore, xShell)

        qext = (2 / xShell2) * np.sum(n1 * (an.real + bn.real))
        qsca = (2 / xShell2) * np.sum(n1 * (an.real ** 2 + an.imag ** 2 + bn.real ** 2 + bn.imag ** 2))
        qabs = qext - qsca

        g1 = [an.real[1:int(nmax)], an.imag[1:int(nmax)], bn.real[1:int(nmax)], bn.imag[1:int(nmax)]]
        g1 = [np.append(x, 0.0) for x in g1]
        g = (4 / (qsca * xShell2)) * np.sum(
            (n2 * (an.real * g1[0] + an.imag * g1[1] + bn.real * g1[2] + bn.imag * g1[3])) + (
                    n3 * (an.real * bn.real + an.imag * bn.imag)))

        qpr = qext - qsca * g
        qback = (1 / xShell2) * (np.abs(np.sum(n1 * ((-1) ** n) * (an - bn))) ** 2)
        qratio = qback / qsca

        if asCrossSection:
            css = np.pi * (dShell / 2) ** 2
            cext = css * qext
            csca = css * qsca
            cabs = css * qabs
            cpr = css * qpr
            cback = css * qback
            cratio = qratio
            if asDict:
                return dict(Cext=cext, Csca=csca, Cabs=cabs, g=g, Cpr=cpr, Cback=cback, Cratio=cratio)
            else:
                return cext, csca, cabs, g, cpr, cback, cratio
        else:
            if asDict:
                return dict(Qext=qext, Qsca=qsca, Qabs=qabs, g=g, Qpr=qpr, Qback=qback, Qratio=qratio)
            else:
                return qext, qsca, qabs, g, qpr, qback, qratio


def CoreShell_ab(mCore, mShell, xCore, xShell):
    m = mShell / mCore
    u = mCore * xCore
    v = mShell * xCore
    w = mShell * xShell

    mx = max(np.abs(mCore * xShell), np.abs(mShell * xShell))
    nmax = np.round(2 + xShell + 4 * (xShell ** (1 / 3)))
    nmx = np.round(max(nmax, mx) + 16)
    n = np.arange(1, nmax + 1)
    nu = n + 0.5

    sv = np.sqrt(0.5 * np.pi * v)
    sw = np.sqrt(0.5 * np.pi * w)
    sy = np.sqrt(0.5 * np.pi * xShell)

    pv = sv * jv(nu, v)
    pw = sw * jv(nu, w)
    py = sy * jv(nu, xShell)

    chv = -sv * yv(nu, v)
    chw = -sw * yv(nu, w)
    chy = -sy * yv(nu, xShell)

    p1y = np.append([np.sin(xShell)], [py[0:int(nmax) - 1]])
    ch1y = np.append([np.cos(xShell)], [chy[0:int(nmax) - 1]])
    gsy = py - (0 + 1.0j) * chy
    gs1y = p1y - (0 + 1.0j) * ch1y

    Dnu = np.zeros((int(nmx)), dtype=complex)
    Dnv = np.zeros((int(nmx)), dtype=complex)
    Dnw = np.zeros((int(nmx)), dtype=complex)
    for i in range(int(nmx) - 1, 1, -1):
        Dnu[i - 1] = i / u - 1 / (Dnu[i] + i / u)
        Dnv[i - 1] = i / v - 1 / (Dnv[i] + i / v)
        Dnw[i - 1] = i / w - 1 / (Dnw[i] + i / w)

    Du = Dnu[1:int(nmax) + 1]
    Dv = Dnv[1:int(nmax) + 1]
    Dw = Dnw[1:int(nmax) + 1]

    uu = m * Du - Dv
    vv = Du / m - Dv
    fv = pv / chv

    dns = ((uu * fv / pw) / (uu * (pw - chw * fv) + (pw / pv) / chv)) + Dw
    gns = ((vv * fv / pw) / (vv * (pw - chw * fv) + (pw / pv) / chv)) + Dw
    a1 = dns / mShell + n / xShell
    b1 = mShell * gns + n / xShell

    an = (py * a1 - p1y) / (gsy * a1 - gs1y)
    bn = (py * b1 - p1y) / (gsy * b1 - gs1y)

    return an, bn


if __name__ == "__main__":
    # MieQ(m, wavelength, diameter, nMedium=1.0, asDict=False, asCrossSection=False)
    # 单球Mie散射
    print(MieQ(1.77 + 0.63j, 375, 300, nMedium=1.33, asDict=True))

    # MieQCoreShell(mCore, mShell, wavelength, dCore, dShell, nMedium=1.0, asDict=False, asCrossSection=False)
    # 双层球Mie散射
    print(MieQCoreShell(1.77 + 0.63j, 1.52 + 0.5j, 375, 300, 500, nMedium=1.33, asDict=True))
